import os
import time
import json
from flask import Flask, render_template, jsonify, send_from_directory, request, Response

app = Flask(__name__)
CONTENT_DIR = os.path.join(os.path.dirname(__file__), 'content')
IMAGE_DIR = os.path.join(os.path.dirname(__file__), 'image')
ZIP_DIR = os.path.join(os.path.dirname(__file__), 'zip')

# 限流设置
REQUEST_LIMIT = 600  # 每小时最多请求次数
TRAFFIC_LIMIT = 300 * 1024 * 1024  # 每小时最多流量（300MB）

# 限流数据结构：{ip: {count: int, bytes: int, timestamp: float}}
rate_limit_data = {}


def clear_old_records():
    """清除一小时前的记录"""
    now = time.time()
    for ip in list(rate_limit_data.keys()):
        if now - rate_limit_data[ip]['timestamp'] > 3600:
            del rate_limit_data[ip]


def custom_json_response(data, status=200):
    return Response(
        json.dumps(data, ensure_ascii=False, indent=None, separators=(',', ':')),
        mimetype='application/json',
        status=status
    )


@app.before_request
def rate_limit():
    clear_old_records()
    ip = request.remote_addr

    if request.path.startswith('/content/') or request.path.startswith('/image/'):
        # 初始化IP记录
        if ip not in rate_limit_data:
            rate_limit_data[ip] = {'count': 0, 'bytes': 0, 'timestamp': time.time()}

        record = rate_limit_data[ip]
        record['count'] += 1

        # 预估响应大小（实际可更精确，此处为简化处理）
        # 这里可以优化为记录实际发送字节数
        if request.path.startswith('/content/'):
            filename = request.path[len('/content/'):]
            filepath = os.path.join(CONTENT_DIR, filename)
            if os.path.exists(filepath):
                size = os.path.getsize(filepath)
            else:
                size = 0
        elif request.path.startswith('/image/'):
            filename = request.path[len('/image/'):]
            filepath = os.path.join(IMAGE_DIR, filename)
            if os.path.exists(filepath):
                size = os.path.getsize(filepath)
            else:
                size = 0
        elif request.path.startswith('/zip/'):
            filename = request.path[len('/zip/'):]
            filepath = os.path.join(ZIP_DIR, filename)
            if os.path.exists(filepath):
                size = os.path.getsize(filepath)
            else:
                size = 0
        else:
            size = 0

        record['bytes'] += size

        if record['count'] > REQUEST_LIMIT:
            return custom_json_response({"error": "请求次数超过每小时限制"}, status=429)

        if record['bytes'] > TRAFFIC_LIMIT:
            return custom_json_response({"error": "流量超过每小时限制"}, status=429)


# 以下为原有路由逻辑，保持不变
def build_nav_tree(path, prefix=""):
    tree = []
    for name in sorted(os.listdir(path)):
        full_path = os.path.join(path, name)
        display_name = name.replace(".md", "").capitalize()
        if os.path.isdir(full_path):
            children = build_nav_tree(full_path, prefix + name + "/")
            tree.append({
                "name": display_name,
                "path": prefix + name,
                "children": children
            })
        elif name.endswith(".md"):
            tree.append({
                "name": display_name,
                "path": prefix + name
            })
    return tree


@app.route('/')
def serve_index():
    return send_from_directory('./webui', 'index.html')


@app.route('/api/nav-tree')
def get_nav_tree():
    tree = build_nav_tree(CONTENT_DIR)
    return jsonify(tree)


@app.route('/content/<path:filename>')
def get_content(filename):
    filepath = os.path.join(CONTENT_DIR, filename)
    if os.path.exists(filepath):
        with open(filepath, 'r', encoding='utf-8') as f:
            return f.read()
    return "File not found", 404


@app.route('/style/<path:filename>')
def send_static(filename):
    return send_from_directory('./webui/style', filename)


# ✅ 添加 image 资源代理访问
@app.route('/image/<path:filename>')
def send_image(filename):
    return send_from_directory(IMAGE_DIR, filename)


# ✅ 添加 zip 资源代理访问
@app.route('/zip/<path:filename>')
def send_zip(filename):
    return send_from_directory(ZIP_DIR, filename)


if __name__ == '__main__':
    app.run(debug=False, host='0.0.0.0', port=5000)
